import numpy as np
import torch 
import torch.nn as nn 
from monai.networks.nets.basic_unet import BasicUNet
from monai.networks.nets.unetr import UNETR
from monai.networks.nets.swin_unetr import SwinUNETR
from monai.inferers import SlidingWindowInferer
from light_training.evaluation.metric import dice
from light_training.trainer import Trainer
from monai.utils import set_determinism
from light_training.utils.lr_scheduler import LinearWarmupCosineAnnealingLR
from light_training.utils.files_helper import save_new_model_and_delete_last
from models.uent2d import UNet2D
from models.uent3d import UNet3D
from models.uent25d import UNet25D
from monai.networks.nets.segresnet import SegResNet
from models.transbts.TransBTS_downsample8x_skipconnection import TransBTS
from models.nestedformer.nested_former import NestedFormer
from models.swinunet2d.swinunet import SwinUnet
from einops import rearrange
from monai.networks.nets.vnet import VNet
from models.modelgenesis.unet3d import UNet3DModelGen
from models.transvw.models.ynet3d import UNet3DTransVW
from monai.networks.nets.attentionunet import AttentionUnet
import random
import h5py
from torch.utils.data import Dataset
from models.hdensenet_huanhu.net import dense_rnn_net
from models.unet_nested.unet_nested_3d import UNet_Nested3D
import glob 
import os 
import SimpleITK as sitk 
import matplotlib.pyplot as plt 

data_paths = glob.glob("./huanhu_raw_data_resampled/label/*/*")
print(len(data_paths))
# data = data_dir[23]
# data = "./huanhu_raw_data_resampled/label/Grade_1/BAI_YUE_BIN"
# t1 = sitk.ReadImage(os.path.join(data, "t1.nii.gz"))
# t1_arr = sitk.GetArrayFromImage(t1)
# t2 = sitk.ReadImage(os.path.join(data, "t2.nii.gz"))
# t2_arr = sitk.GetArrayFromImage(t2)
# label = sitk.ReadImage(os.path.join(data, "tumor.nii.gz"))
# label_arr = sitk.GetArrayFromImage(label)

# for i in range(20):
#     plt.subplot(1, 3, 1)
#     plt.imshow(t1_arr[i], cmap="gray")
#     plt.subplot(1, 3, 2)
#     plt.imshow(t2_arr[i], cmap="gray")
#     plt.subplot(1, 3, 3)
#     plt.imshow(label_arr[i], cmap="gray")
#     plt.show()
# exit(0)
torch.backends.cudnn.deterministic = True
random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
logdir = "./logs_huanhu_tumor_only/segresnet"

env = "pytorch"
model_save_path = os.path.join(logdir, "model")
max_epoch = 300 ## att-unet

batch_size = 1
val_every = 5
num_gpus = 1
device = "cuda:0"

train_two_fold = []
test_two_fold = []
train_paths = []
test_paths = []

for i in range(120):
    train_paths.append(data_paths[i])
#
for i in range(120, len(data_paths)):
    test_paths.append(data_paths[i])

train_two_fold.append(train_paths)
test_two_fold.append(test_paths)

# ## 第二折
train_paths = []
test_paths = []
for i in range(120, len(data_paths)):
    train_paths.append(data_paths[i])

for i in range(120):
    test_paths.append(data_paths[i])

train_two_fold.append(train_paths)
test_two_fold.append(test_paths)


class Dataset3d(Dataset):
    """
    """
    def __init__(
        self,
        paths
    ) -> None:

        super(Dataset3d, self).__init__()
        self.cache_num = len(paths)
        if self.cache_num > 0:
            self._cache_image = [None] * self.cache_num

            self._cache_label = [None] * self.cache_num

            for i in range(self.cache_num):
                self._cache_image[i], self._cache_label[i] = \
                        self._load_cache_item(paths[i])

    def get_labels(self, label):
        labels = np.zeros(label.shape[1:])
        labels[label[0] == 1] = 1
        # labels[label[1] == 1] = 2
        return labels

    def _load_cache_item(self, d_path):
        t1 = sitk.ReadImage(os.path.join(d_path, "t1.nii.gz"))
        t1_arr = sitk.GetArrayFromImage(t1)

        t2 = sitk.ReadImage(os.path.join(d_path, "t2.nii.gz"))
        t2_arr = sitk.GetArrayFromImage(t2)

        label = sitk.ReadImage(os.path.join(d_path, "tumor.nii.gz"))
        label_arr = sitk.GetArrayFromImage(label)

        arr = np.stack([t1_arr, t2_arr], axis=0).astype(np.float32)
        # labels = self.get_labels(label)

        return arr, label_arr

    def __getitem__(self, index):
        image = self._cache_image[index]
        image_mean = image.mean()
        image_std = image.std()
        image = (image - image_mean) / image_std

        label = self._cache_label[index]

        return image.astype(np.float32), label.astype(np.float32)

    def __len__(self):
        return len(self._cache_image)

class FuseUNet(nn.Module):
    def __init__(self, in_ch, out_ch) -> None:
        super().__init__()
        self.model3d = UNet3D(in_ch, out_ch)
        self.model25d = UNet25D(in_ch, out_ch)
        self.model2d = UNet2D(in_ch, out_ch)
       

    def forward(self, x, pred_type="all"):
        if pred_type == "all":
            pred_3d = self.model3d(x)
            pred_25d = self.model25d(x)
            pred_2d = self.model2d(x)

            return pred_3d  + pred_25d + pred_2d
            # return pred_3d, pred_25d, pred_2d
        
        elif pred_type == "3d":
            return self.model3d(x)
        
        elif pred_type == "25d":
            return self.model25d(x)

        elif pred_type == "2d":
            return self.model2d(x)


class BraTSTrainer(Trainer):
    def __init__(self, env_type, max_epochs, batch_size, device="cpu", val_every=1, num_gpus=1, logdir="./logs/", master_ip='localhost', master_port=17750, training_script="train.py",fold=0):
        super().__init__(env_type, max_epochs, batch_size, device, val_every, num_gpus, logdir, master_ip, master_port, training_script)
    
        self.fold=fold  
        self.model_name = "1"
        # self.model = SwinUNETR([96, 96, 96], 4, 4)
        # self.model = UNet2D()
        # self.model = UNet3D()
        # self.model = UNETR(4, 4, [96, 96, 96])
        # _, model = TransBTS(dataset='huanhu', _conv_repr=True, _pe_type="learned")
        # self.model = model
        # self.model = FuseUNet(in_ch=1, out_ch=2)

        self.model = SegResNet(3, 16, 2, 2)

        # self.model = VNet(3, 1, 2, bias=True)
        # self.model_name = "hdensenet"
        # self.model = dense_rnn_net(32, 2, drop_rate=0.0)

        # self.model_name = "swinunet2d"
        # from models.swinunet2d.config import get_config
        # config = get_config("/home/xingzhaohu/jiuding_code/mutual_learning_huanhu/models/swinunet2d/swin_tiny_patch4_window7_224_lite.yaml")
        # self.model = SwinUnet(config, img_size=224, num_classes=2, in_channels=3)
        # self.model.load_from(config)

        # self.model = SwinUNETR2D()
        # self.model = UNet_Nested3D(in_channels=1, n_classes=2, is_deconv=False)
        # self.model = UNet3DModelGen(4)
        # self.model = UNet3DTransVW(4)
        # weight_dir = "/home/xingzhaohu/jiuding_code/mutual_learning/logs_brats/Genesis_Chest_CT.pt"
        # weight_dir = "/home/xingzhaohu/jiuding_code/mutual_learning/logs_brats/TransVW_chest_ct.pt"
        # checkpoint = torch.load(weight_dir)
        # state_dict = checkpoint['state_dict']
        # unParalled_state_dict = {}
        # for key in state_dict.keys():
        #     if "down_tr64.ops.0.conv1.weight" in key:
        #         state_dict[key] = state_dict[key].repeat(1, 4, 1, 1, 1)

        #     if "out_tr" not in key:
        #         unParalled_state_dict[key.replace("module.", "")] = state_dict[key]

        # self.model.load_state_dict(unParalled_state_dict, strict=False)

        # self.model = AttentionUnet(3, 1, 2, channels=[16, 32, 64, 128], strides=[2, 2, 2, 2])

        self.best_mean_dice = 0.0
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-3, weight_decay=1e-5)
        # self.scheduler = LinearWarmupCosineAnnealingLR(self.optimizer,
        #                                           warmup_epochs=30,
        #                                           max_epochs=max_epochs,
        #                                           warmup_start_lr=1e-6)
        self.scheduler = LinearWarmupCosineAnnealingLR(self.optimizer,
                                                  warmup_epochs=30,
                                                  max_epochs=max_epochs,
                                                    )
        self.loss_func = nn.CrossEntropyLoss()

    def training_step(self, batch):
        import time 
        image, label = self.get_input(batch)

        pred = self.model(image)
        loss = self.loss_func(pred, label.long())
        self.log("train_loss", loss, step=self.global_step)
        return loss 

    # for image, label in data_loader:
    def get_input(self, batch):
        image, label = batch 
        # if self.model_name != "swinunet2d":
        #     image = nn.functional.interpolate(image, size=(24, 256, 256), mode="trilinear", align_corners=False)
        #     label = torch.unsqueeze(label, dim=1)
        #     label = nn.functional.interpolate(label, size=(24, 256, 256), mode="nearest")
        #     label = torch.squeeze(label, dim=1).long()
        # else :
        #     image = nn.functional.interpolate(image, size=(32, 224, 224), mode="trilinear", align_corners=False)
        #     label = torch.unsqueeze(label, dim=1)
        #     label = nn.functional.interpolate(label, size=(32, 224, 224), mode="nearest")
        #     label = torch.squeeze(label, dim=1).long()

        return image, label 

    def validation_step(self, batch):
        image, label = self.get_input(batch)
       
        output = self.model(image).argmax(dim=1).cpu().numpy()
        target = label.cpu().numpy()
        o = output == 1; t = target == 1 # ce
        dice_score_tumor = dice(o, t)

        # o = output == 2; t = target == 2 # ce
        # dice_score_edma = dice(o, t)

        return dice_score_tumor

    def validation_end(self, mean_val_outputs, val_outputs=None):
        dice_score_tumor = mean_val_outputs

        self.log("dice_score_tumor", dice_score_tumor, step=self.epoch)
        # self.log("dice_score_edma", dice_score_edma, step=self.epoch)
        # dice_score = (dice_score_tumor + dice_score_edma) / 2

        if dice_score_tumor > self.best_mean_dice:
            self.best_mean_dice = dice_score_tumor
            save_new_model_and_delete_last(self.model, 
                                            os.path.join(model_save_path, 
                                            f"best_model_{dice_score_tumor:.4f}_{self.fold}.pt"), 
                                            delete_symbol=f"best_model_*_{self.fold}")

        save_new_model_and_delete_last(self.model, 
                                        os.path.join(model_save_path, 
                                        f"final_model_{dice_score_tumor:.4f}_{self.fold}.pt"), 
                                        delete_symbol=f"final_model_*_{self.fold}")

        print(f"dice_score is {dice_score_tumor}")

if __name__ == "__main__":

    fold = 0

    trainer = BraTSTrainer(env_type=env,
                            max_epochs=max_epoch,
                            batch_size=batch_size,
                            device=device,
                            logdir=logdir,
                            val_every=val_every,
                            num_gpus=num_gpus,
                            master_port=17751,
                            training_script=__file__,
                            fold=fold,
                            )

    train_paths = train_two_fold[fold]
    test_paths = test_two_fold[fold]
    train_ds = Dataset3d(train_paths)
    val_ds = Dataset3d(test_paths)


    trainer.train(train_dataset=train_ds, val_dataset=val_ds)

    print(f"第{fold}折训练完成。。。。")
